# ======================================================================
# Copyright CERFACS (January 2019)
# Contributor: Adrien Suau (adrien.suau@cerfacs.fr)
#
# This software is governed by the CeCILL-B license under French law and
# abiding  by the  rules of  distribution of free software. You can use,
# modify  and/or  redistribute  the  software  under  the  terms  of the
# CeCILL-B license as circulated by CEA, CNRS and INRIA at the following
# URL "http://www.cecill.info".
#
# As a counterpart to the access to  the source code and rights to copy,
# modify and  redistribute granted  by the  license, users  are provided
# only with a limited warranty and  the software's author, the holder of
# the economic rights,  and the  successive licensors  have only limited
# liability.
#
# In this respect, the user's attention is drawn to the risks associated
# with loading,  using, modifying and/or  developing or reproducing  the
# software by the user in light of its specific status of free software,
# that  may mean  that it  is complicated  to manipulate,  and that also
# therefore  means that  it is reserved for  developers and  experienced
# professionals having in-depth  computer knowledge. Users are therefore
# encouraged  to load and  test  the software's  suitability as  regards
# their  requirements  in  conditions  enabling  the  security  of their
# systems  and/or  data to be  ensured and,  more generally,  to use and
# operate it in the same conditions as regards security.
#
# The fact that you  are presently reading this  means that you have had
# knowledge of the CeCILL-B license and that you accept its terms.
# ======================================================================

"""Module implementing Hamiltonian simulation for 1-sparse integer Hamiltonians.

This module contains the functions used to simulate a 1-sparse Hamiltonian that contain
either only real integer entries or only imaginary integer entries.

The circuit returned by these procedures simulate exactly the given Hamiltonian.
"""

from qaths.simulation.base._routines import A, G, exp_ZFt, exp_ZZFt
from qat.lang.AQASM.routines import QRoutine
from qat.lang.AQASM.misc import build_gate
from qat.lang.AQASM.gates import AbstractGate


@build_gate(
    "simulate_unsigned_integer_weighted_hamiltonian",
    [str, int, int, float],
    arity=lambda name, n, i, t: n,
)
def simulate_unsigned_integer_weighted_hamiltonian(
    oracle_name: str, n: int, int_size: int, time: float
) -> QRoutine:
    r"""Simulate the evolution under an Hamiltonian with at most one positive integer \
    entry in each row/column.

    Let :math:`n` be the size of :math:`\ket{x}`. The returned routine needs :math:`n`
    qubits organised as follow: ::

        |        |x>        |
        |   .   .   .   .   |
        |   0          n-1  |

    * x is the starting state of the evolution
      :math:`\left(\ket{x} = \ket{\phi_0}\right)`.

    :param oracle_name: name of the oracle that encodes the Hamiltonian matrix we
        want to simulate.
        The name will be used to create an AbstractGate, that can be linked afterwards.

        The oracle :math:`O` take as input :math:`\ket{x}\ket{0}\ket{0}`
        organised as follow: ::

            |        |x>        |        |m>        |        |v>        |
            |   .   .   .   .   |   .   .   .   .   |   .   .   .   .   |
            |  0           n-1  |   n         2*n-1 |  2*n       2*n+i-1|

        and outputs the quantum state
        :math:`\ket{x}\ket{m(x)}\ket{v(x)}` with:

        #. :math:`\ket{x}` encoding the index of the considered row (all the indices are
           in superposition).
        #. :math:`\ket{m(x)}` encoding the index of the column of the only non-zero
           element in the column of index :math:`x`.
        #. :math:`\ket{v(x)}` encoding the absolute value of the only non-zero entry
           in the column of index :math:`x`.

    :param n: The number of qubits the simulated Hamiltonian acts on.
    :param int_size: The number of qubit used to represent an integer.
    :param time: The duration of the desired evolution.
    """

    routine = QRoutine()
    # Aliases to make the code more readable.
    x = routine.new_wires(n)
    m = routine.new_wires(n)
    v = routine.new_wires(int_size)
    p = routine.new_wires(1)
    routine.set_ancillae(m)
    routine.set_ancillae(v)
    routine.set_ancillae(p)

    oracle = AbstractGate(oracle_name, [], arity=2 * n + int_size)

    routine.protected_apply(oracle(), x, m, v)
    routine.protected_apply(A(n), x, m, p)

    routine.apply(exp_ZFt(int_size, time), p, v)

    routine.protected_apply(A(n).dag(), x, m, p)
    routine.protected_apply(oracle().dag(), x, m, v)

    return routine


@build_gate(
    "simulate_signed_integer_weighted_hamiltonian",
    [str, int, int, float],
    arity=lambda name, n, i, t: n,
)
def simulate_signed_integer_weighted_hamiltonian(
    oracle_name: str, n: int, int_size: int, time: float
) -> QRoutine:
    r"""Simulate the evolution under an Hamiltonian with at most one integer entry in \
    each row/column.

    Let :math:`n` be the size of :math:`\ket{x}`. The returned routine needs :math:`n`
    qubits organised as follow: ::

        |    |x>    |
        |   .   .   |
        |   0  n-1  |

    * x is the starting state of the evolution
      :math:`\left(\ket{x} = \ket{\phi_0}\right)`.

    :param oracle_name: name of the oracle that encodes the Hamiltonian matrix we
        want to simulate.
        The name will be used to create an AbstractGate, that can be linked afterwards.

        The oracle :math:`O` take as input :math:`\ket{x}\ket{0}\ket{0}\ket{0}`
        organised as follow: ::

            |        |x>        |       |m>        |       |v>        |  |s>  |
            |   .   .   .   .   |  .   .   .   .   |  .   .   .   .   |   .   |
            |  0           n-1  |  n         2*n-1 | 2*n        2n+i-1|  2n+i |

        and outputs the quantum state
        :math:`\ket{x}\ket{m(x)}\ket{v(x)}\ket{s(x)}` with:

        #. :math:`\ket{x}` encoding the index of the considered row (all the indices are
           in superposition).
        #. :math:`\ket{m(x)}` encoding the index of the column of the only non-zero
           element in the column of index :math:`x`.
        #. :math:`\ket{v(x)}` encoding the absolute value of the only non-zero entry
           in the column of index :math:`x`.
        #. :math:`\ket{s(x)}` encoding the sign of the only non-zero entry in the column
           of index :math:`x` with the convention

           .. math::

              s(x) = \left\{\begin{array}{lr}
              0 & \text{if } x \geqslant 0 \\
              1 & \text{else} \\
              \end{array}\right.

    :param n: The number of qubits the simulated Hamiltonian acts on.
    :param int_size: The number of qubit used to represent an integer.
    :param time: The duration of the desired evolution.
    """
    routine = QRoutine()
    # Aliases to make the code more readable.
    x = routine.new_wires(n)
    m = routine.new_wires(n)
    w = routine.new_wires(int_size)
    p = routine.new_wires(1)
    s = routine.new_wires(1)
    routine.set_ancillae(m)
    routine.set_ancillae(w)
    routine.set_ancillae(p)
    routine.set_ancillae(s)

    oracle = AbstractGate(oracle_name, [], arity=2 * n + int_size + 1)

    routine.protected_apply(oracle(), x, m, w, s)
    routine.protected_apply(A(n), x, m, p)

    routine.apply(exp_ZZFt(int_size, time), p, s, w)

    routine.protected_apply(A(n).dag(), x, m, p)
    routine.protected_apply(oracle().dag(), x, m, w, s)

    return routine


@build_gate(
    "simulate_imaginary_integer_weighted_hamiltonian",
    [str, int, int, float],
    arity=lambda name, n, i, t: n,
)
def simulate_imaginary_integer_weighted_hamiltonian(
    oracle_name: str, n: int, int_size: int, time: float
) -> QRoutine:
    r"""Simulate the evolution under an Hamiltonian with at most one imaginary integer \
    entry in each row/column.

    Let :math:`n` be the size of :math:`\ket{x}`. The returned routine needs :math:`n`
    qubits organised as follow: ::

        |    |x>    |
        |   .   .   |
        |   0  n-1  |

    * x is the starting state of the evolution
      :math:`\left(\ket{x} = \ket{\phi_0}\right)`.

    :param oracle_name: name of the oracle that encodes the Hamiltonian matrix we
        want to simulate.
        The name will be used to create an AbstractGate, that can be linked afterwards.

        The oracle :math:`O` take as input :math:`\ket{x}\ket{0}\ket{0}\ket{0}`
        organised as follow: ::

            |        |x>        |       |m>        |       |v>        |  |s>  |
            |   .   .   .   .   |  .   .   .   .   |  .   .   .   .   |   .   |
            |  0           n-1  |  n         2*n-1 | 2*n        2n+i-1|  2n+i |

        and outputs the quantum state
        :math:`\ket{x}\ket{m(x)}\ket{v(x)}\ket{s(x)}` with:

        #. :math:`\ket{x}` encoding the index of the considered row (all the indices are
           in superposition).
        #. :math:`\ket{m(x)}` encoding the index of the column of the only non-zero
           element in the column of index :math:`x`.
        #. :math:`\ket{v(x)}` encoding the absolute value of the only non-zero entry
           in the column of index :math:`x`.
        #. :math:`\ket{s(x)}` encoding the sign of the only non-zero entry in the column
           of index :math:`x` with the convention

           .. math::

              s(x) = \left\{\begin{array}{lr}
              0 & \text{if } x \geqslant 0 \\
              1 & \text{else} \\
              \end{array}\right.

    :param n: The number of qubits the simulated Hamiltonian acts on.
    :param int_size: The number of qubits used to represent an integer.
    :param time: The duration of the desired evolution.
    """
    routine = QRoutine()
    # Aliases to make the code more readable.
    x = routine.new_wires(n)
    m = routine.new_wires(n)
    w = routine.new_wires(int_size)
    p = routine.new_wires(1)
    s = routine.new_wires(1)
    routine.set_ancillae(m)
    routine.set_ancillae(w)
    routine.set_ancillae(p)
    routine.set_ancillae(s)

    oracle = AbstractGate(oracle_name, [], arity=2 * n + int_size + 1)

    # Prepare the quantum state
    routine.protected_apply(oracle(), x, m, w, s)
    routine.protected_apply(A(n), x, m, p)
    routine.protected_apply(G(), s)

    # Perform the simulation
    # TODO: explain why we need a minus here.
    # Linked with the problem of negative sign representation (in the paper it's
    # positive -> |0>, negative -> |1>, in practice only the opposite works).
    routine.apply(exp_ZZFt(int_size, -time), p, s, w)

    # Inverse the preparatory step.
    routine.protected_apply(G().dag(), s)
    routine.protected_apply(A(n).dag(), x, m, p)
    routine.protected_apply(oracle().dag(), x, m, w, s)

    return routine
